#include "memory.h"
#include "print.h"
#include "debug.h"
#include "string.h"
#include "thread.h"
#include "sync.h"

#define PG_SIZE             4096

#define PDE_IDX(addr) ((addr & 0xffc00000) >> 22)
#define PTE_IDX(addr) ((addr & 0x003ff000) >> 12)

/* 0xc009f000是内核线程主栈顶 0xc009e000是内核PCB所在位置
 * 一个页表大小即4K的位图可管理128M内存
 * 我们预留4个页表大小 最大管理512M内存 方便拓展
 * 所以位图起始地址应该是 0xc009e000 - 4K*4 = 0xc009a000
 */
#define MEM_BITMAP_BASE 0xc009a000

/* 0xc0000000是内核从虚拟地址3G起. 0x100000意指跨过低端1M内存,使虚拟地址在逻辑上连续 */
#define K_HEAP_START 0xc0100000

struct mem_pool{
    struct lock lock;
    struct bitmap pbtmp;
    uint32_t pool_start_addr;
    uint32_t pool_size;
};
struct mem_pool kernel_pool, user_pool;
struct virtual_addr_pool kernel_vaddr_pool;

static void* get_vaddr(enum pool_flags pf, uint32_t pg_cnt){
    int _idx_bit = -1;
    uint32_t _cnt = 0, _vaddr_start;
    if(pf == PF_KERNEL){
        /* 查找pg_cnt个可用的位 */
        _idx_bit = bitmap_scan(&kernel_vaddr_pool.vbtmp, pg_cnt);
        if(_idx_bit == -1){
            return NULL;
        }
        /* 将刚才查找到的位置为1 即标记为已用 */
        while(_cnt < pg_cnt){
            bitmap_set(&kernel_vaddr_pool.vbtmp, _idx_bit + _cnt++, 1);
        }
        _vaddr_start = kernel_vaddr_pool.vaddr_start + _idx_bit * PG_SIZE;
    }else{
        struct task_struct* __cur = running_task();

        _idx_bit = bitmap_scan(&__cur->userprog_vaddr.vbtmp, pg_cnt);
        if(_idx_bit == -1){
            return NULL;
        }

        while(_cnt < pg_cnt){
            bitmap_set(&__cur->userprog_vaddr.vbtmp, _idx_bit + _cnt++, 1);
        }
        _vaddr_start = __cur->userprog_vaddr.vaddr_start + _idx_bit * PG_SIZE;
        /* (0xc0000000 - PG_SIZE)做为用户3级栈已经在start_process被分配 
         * 而且0xc0000000以上是内核的地址
        */
        ASSERT((uint32_t)_vaddr_start < (0xc0000000 - PG_SIZE));
    }
    return (void*)_vaddr_start;
}

static void* get_paddr(struct mem_pool* pool){
    /* 位图中查找1个可用的位 */
    int _idx_bit = bitmap_scan(&pool->pbtmp, 1);
    if(_idx_bit == -1){
        return NULL;
    }
    /* 将刚才查找到的位置为1 即标记为已用 */
    bitmap_set(&pool->pbtmp, _idx_bit, 1);

    uint32_t phy_addr = (pool->pool_start_addr + (_idx_bit * PG_SIZE));

    return (void*)phy_addr;
}
/* 将虚拟地址和物理地址做映射 */
static void page_table_map(void* vaddr_, void* paddr_){
    uint32_t vaddr = (uint32_t) vaddr_, paddr = (uint32_t)paddr_;
    uint32_t* pde = pde_ptr(vaddr);
    uint32_t* pte = pte_ptr(vaddr);

    /* 如果页表存在，则不需要创建页表 */
    if(*pde & PG_P_1){
        /* 防止pte已经存在 */
        ASSERT(!(*pte & PG_P_1));
        *pte = paddr | PG_US_U | PG_RW_W | PG_P_1;
    }else{
        /* 如果页表不存在，则需要创建页表，额外申请一页大小创建页表 */
        uint32_t pt_paddr = (uint32_t)get_paddr(&kernel_pool);
        /* 将pde设置为这个页表的地址 */
        *pde = pt_paddr | PG_US_U | PG_RW_W | PG_P_1;
        /* 由于这个页表的4k大小是刚申请的，所以我们要给他清零
         * pte是页表中某个项的地址，后12位（明确地说是只用了10位）用来索引
         * 哪个entry,所以将后12位清零就可以找到页表第一个entry的地址
         */
        memset((void*)((uint32_t)pte & 0xfffff000), 0, PG_SIZE);
        *pte = paddr | PG_US_U | PG_RW_W | PG_P_1;
    }   
}

uint32_t* pte_ptr(uint32_t vaddr){
    uint32_t* pte = (uint32_t*)((0xffc00000) + ((vaddr & 0xffc00000) >> 10) + 
                                (PTE_IDX(vaddr) * 4));
    return pte;
}

uint32_t* pde_ptr(uint32_t vaddr){
    uint32_t* pde = (uint32_t*)((0xfffff000) + (PDE_IDX(vaddr) * 4));

    return pde;
}

void* malloc_page(enum pool_flags pf, uint32_t pg_cnt){
    struct mem_pool* _mem_pool = pf == PF_KERNEL ? &kernel_pool : &user_pool;
    /* 保证所申请的页数不大于位图所能管理的页数 */
    ASSERT(pg_cnt > 0 && pg_cnt < _mem_pool->pbtmp.btmp_bytes_len);

    /* 获取pg_cnt个页的虚拟地址 */
    void* _vaddr_start = get_vaddr(pf, pg_cnt);
    if(_vaddr_start == NULL){
        return NULL;
    }

    uint32_t _vaddr = (uint32_t)_vaddr_start;
    while(pg_cnt-- > 0){
        /* 因为虚拟地址是连续的，但物理地址可以不连续
         * 所以虚拟地址直接连续申请，物理地址一个一个申请
         * 逐页做映射，体现了乱序映射的魅力
         */
        void* __paddr = get_paddr(_mem_pool);
        if(__paddr == NULL){
            put_str("get_paddr error\n");
            //todo-jianglin: memory rollback
            return NULL;
        }
        page_table_map((void*)_vaddr, __paddr);
        _vaddr += PG_SIZE;
    }
    
    return _vaddr_start;
}

void* get_kernel_pages(uint32_t pg_cnt){
    void* vaddr =  malloc_page(PF_KERNEL, pg_cnt);
    if(vaddr != NULL){
        memset(vaddr, 0, pg_cnt * PG_SIZE);
    }
    return vaddr;
}

void* get_user_pages(uint32_t pg_cnt){
    lock_acquire(&user_pool.lock);
    void* vaddr =  malloc_page(PF_USER, pg_cnt);
    if(vaddr != NULL){
        memset(vaddr, 0, pg_cnt * PG_SIZE);
    }
    lock_release(&user_pool.lock);
    return vaddr;
}

void* get_a_page(enum pool_flags pf, uint32_t vaddr) {
    struct mem_pool* mem_pool = pf & PF_KERNEL ? &kernel_pool : &user_pool;
    lock_acquire(&mem_pool->lock);
    struct task_struct* cur = running_task();
    int32_t bit_idx = -1;
    
    if(pf == PF_USER && cur->pgdir != NULL){
        /* 内核进程的虚拟地址则在进程自己的虚拟地址池中分配 */
        bit_idx = (vaddr - cur->userprog_vaddr.vaddr_start) / PG_SIZE;
        ASSERT(bit_idx > 0);
        bitmap_set(&cur->userprog_vaddr.vbtmp, bit_idx, 1);

    }else if(pf == PF_KERNEL && cur->pgdir == NULL){
        /* 内核线程的虚拟地址则在kernel_vaddr_pool中分配 */
        bit_idx = (vaddr - kernel_vaddr_pool.vaddr_start) / PG_SIZE;
        ASSERT(bit_idx > 0);
        bitmap_set(&kernel_vaddr_pool.vbtmp, bit_idx, 1);
    }else{
        PANIC("get_a_page:not allow kernel alloc userspace or user alloc kernelspace by get_a_page");
    }

    void* page_phyaddr = get_paddr(mem_pool);
    if(page_phyaddr == NULL){
        lock_release(&mem_pool->lock);
        return NULL;
    }
    page_table_map((void*)vaddr, page_phyaddr);

    lock_release(&mem_pool->lock);

    return (void*)vaddr;
}

/* 得到虚拟地址映射到的物理地址 */
uint32_t addr_v2p(uint32_t vaddr) {
   uint32_t* pte = pte_ptr(vaddr);
/* (*pte)的值是页表所在的物理页框地址,
 * 去掉其低12位的页表项属性+虚拟地址vaddr的低12位 */
   return ((*pte & 0xfffff000) + (vaddr & 0x00000fff));
}


static void mem_pool_init(uint32_t mem_size){
    put_str("       mem_pool_init start\n");
    /* 在1M地址之后 有1个页目录项， 页目录项第0个和第768个指向同一个页表
     * 页目录项和页表大小都是4K 随后还有254个页表 总共256个页表大小
     */
    uint32_t page_table_size = 256 * PG_SIZE;
    /* 低端1M内存给内核使用 不算入管理的范围 所以已经使用的内存大小是
     * 页表大小+低端1M
     */
    uint32_t used_mem = page_table_size + 0x100000;
    /* 剩余内存字节大小为全部-已用 */
    uint32_t free_mem = mem_size - used_mem;
    /* 计算剩余的页大小 不满足4k即一页大小的丢弃 */
    uint32_t all_free_page = free_mem / PG_SIZE;
    /* 将剩余内存页数量公平分配给内核和用户 */
    uint32_t k_free_page = all_free_page / 2;
    uint32_t u_free_page = all_free_page - k_free_page;
    /* 内核池开始地址是已经使用的内存 */
    uint32_t kp_start_addr = used_mem;
    /* 用户池开始地址是内核池地址+内核池大小 */
    uint32_t up_start_addr = used_mem + k_free_page * PG_SIZE;
    /* 位图每一个位代表一页大小，所以用页数除以8就是位图大小 */
    uint32_t kbtmp_len = k_free_page / 8;
    uint32_t ubtmp_len = u_free_page / 8;

    put_str("       used_mem:");
    put_int_x(used_mem);
    put_str("\n       free_mem:");
    put_int_x(free_mem);
    put_str("\n       all_free_page:");
    put_int_x(all_free_page);
    put_str("\n       k_free_page:");
    put_int_x(k_free_page);
    put_str("\n       u_free_page:");
    put_int_x(u_free_page);
    put_str("\n       kp_start_addr:");
    put_int_x(kp_start_addr);
    put_str("\n       up_start_addr:");
    put_int_x(up_start_addr);
    put_str("\n       kbtmp_len:");
    put_int_x(kbtmp_len);
    put_str("\n       ubtmp_len:");
    put_int_x(ubtmp_len);
    put_char('\n');

    /* 为内核池赋值 */
    kernel_pool.pbtmp.btmp_bytes_len = kbtmp_len;
    kernel_pool.pbtmp.bits = (void*)(MEM_BITMAP_BASE);
    kernel_pool.pool_size = k_free_page * PG_SIZE;
    kernel_pool.pool_start_addr = kp_start_addr;

    /* 为用户池赋值 */
    user_pool.pbtmp.btmp_bytes_len = ubtmp_len;
    user_pool.pbtmp.bits = (void*)(MEM_BITMAP_BASE + kbtmp_len);
    user_pool.pool_size = u_free_page * PG_SIZE;
    user_pool.pool_start_addr = up_start_addr;

    /* 初始化内核虚拟地址池 */
    /* 首先设置内核虚拟地址池中的位图大小与内核物理内存池位图大小一致 */
    kernel_vaddr_pool.vbtmp.btmp_bytes_len = kbtmp_len;
    /* 设置内核虚拟池中的位图开始地址紧跟在用户物理内存池位图后 */
    kernel_vaddr_pool.vbtmp.bits = (void*)(MEM_BITMAP_BASE + kbtmp_len + ubtmp_len);
    /* 0xc0010000以下的地址被映射为内核低端1M字节 所以从0xc0010000开始 */
    kernel_vaddr_pool.vaddr_start = K_HEAP_START;

    bitmap_init(&kernel_pool.pbtmp);
    bitmap_init(&user_pool.pbtmp);
    bitmap_init(&kernel_vaddr_pool.vbtmp);

    put_str("       mem_pool_init done\n");
}

void mem_init(void){
    put_str("   mem_init start\n");
    lock_init(&user_pool.lock);
    lock_init(&kernel_pool.lock);
    uint32_t total_mem_bytes = *(uint32_t*)(0xb00);
    mem_pool_init(total_mem_bytes);
    put_str("   mem_init done\n\n");
}